from datetime import datetime, date
import backtrader as bt
import akshare as ak
import pandas as pd
import numpy as np
import quantstats as qs
import pyfolio as pf
import webbrowser
from sqlalchemy import create_engine, distinct, or_, and_
import pymssql
from urllib.parse import quote_plus as urlquote
from configparser import ConfigParser
from sklearn.linear_model import LinearRegression
import warnings
import random

warnings.filterwarnings('ignore')

pd.set_option('display.max_columns', None)
pd.set_option('display.max_rows', None)

conf = ConfigParser()
conf.read('env.ini')

sqlserver = ('sqlserver_out', 'sqlserver_in')[0]

host = conf.get(sqlserver, 'host')
port = conf.get(sqlserver, 'port')
user = conf.get(sqlserver, 'user')
password = conf.get(sqlserver, 'password')
schema = conf.get(sqlserver, 'schema')
db_name = conf.get(sqlserver, 'db')


def get_stocks():
    sql = "SELECT STK_CODE FROM ETF_INDEX_INFO"
    r = exec_sql(sql)
    df1 = pd.DataFrame(r, columns=['代码'])
    df2 = ak.fund_etf_spot_em()
    res = pd.merge(df1, df2, on=['代码'])
    res = res.loc[res['最新价'] < 10, :]
    stocks = res['代码'].to_list()
    # stocks = [tok for tok in stocks if tok != '512870']

    return stocks


def get_data(stocks, freq='day', source='db'):
    if freq == '1min':
        table_name = 'STOCK_MIN_DATA'
    elif freq == '15min':
        table_name = 'STOCK_15MIN_DATA'
    else:
        table_name = 'STOCK_DAY_DATA'  # HIS_DAILY_DATA
        # table_name = 'ETF_DAY_DATA'

    res = dict()
    date_data = set_date()
    bt = date_data['bt']
    st = date_data['st']
    et = date_data['et']
    bt_d = date_data['bt_d']
    st_d = date_data['st_d']
    et_d = date_data['et_d']

    if source == 'db':
        print('db')
        stocks_str = ','.join(f"'{tok}'" for tok in stocks)
        sql = f"SELECT code, datetime, [close], [open], low, high, volume FROM {table_name} WHERE code in ({stocks_str})"
        r = exec_sql(sql)
        df = pd.DataFrame(r, columns=['code', 'datetime', 'close', 'open', 'low', 'high', 'volume'])
        # df['datetime'] = df['datetime'].map(time_map)

        for code in stocks:
            if code == '399300':
                sig_df = ak.stock_zh_index_daily(symbol="sz399300")
                sig_df = sig_df.loc[(sig_df.date >= st_d) & (sig_df.date <= et_d), :]
                sig_df.rename(columns={'date': 'datetime'}, inplace=True)
                sig_df.insert(0, 'code', code)
                sig_df = sig_df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
                sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])
            else:
                sig_df = df.loc[df['code'] == code]
                sig_df['datetime'] = pd.to_datetime(sig_df['datetime'])

                sig_df = sig_df.loc[(sig_df.datetime >= bt) & (sig_df.datetime <= et), :]

            sig_df.sort_values('datetime', inplace=True, ignore_index=True)
            sig_df.set_index('datetime', drop=False, inplace=True)
            # if len(sig_df) <= 220:
            #     continue
            res[code] = sig_df
    else:
        print('internet')
        for code in stocks:
            if code == '399300':
                df = ak.stock_zh_index_daily(symbol="sz399300")
                df = df.loc[(df.date >= st_d) & (df.date <= et_d), :]
                df.rename(columns={'date': 'datetime'}, inplace=True)
            else:
                df = ak.fund_etf_hist_em(symbol=code, period="daily", start_date=st, end_date=et, adjust="qfq")
                df.rename(columns={'日期': 'datetime', '开盘': 'open', '收盘': 'close', '最高': 'high',
                                   '最低': 'low', '成交量': 'volume'}, inplace=True)

            df.insert(0, 'code', code)
            df = df.loc[:, ['code', 'datetime', 'close', 'open', 'low', 'high', 'volume']]
            df['datetime'] = pd.to_datetime(df['datetime'])
            df.set_index('datetime', inplace=True)
            res[code] = df

    return res


def exec_sql(sql):
    conn = pymssql.connect(host=host, port=port, user=user, password=password, database=db_name)
    cursor = conn.cursor()
    cursor.execute(sql)
    r = cursor.fetchall()
    cursor.close()
    conn.close()
    return r


def time_map(x):
    t = datetime.strptime(str(x), '%Y-%m-%d')
    # t = datetime.strptime(str(x), '%Y%m%d%H%M')
    return t


def insert_db(df, table_name):
    yconnect = create_engine(f'mssql+pymssql://{user}:{urlquote(password)}@{host}:{port}/{db_name}?charset=utf8')
    pd.io.sql.to_sql(df, table_name, yconnect, schema=schema, if_exists='append', index=False)


stocks_info = dict()
init_cash = 3e8


def run_backtrade(stocks):
    cerebro = bt.Cerebro()
    # cerebro.optstrategy(GridReversionStrategy, style=range(1, 4, 1))
    cerebro.addstrategy(GridReversionStrategy)
    data_all = get_data(stocks, source='db')
    index = stocks[0]
    data0 = data_all[index]
    datafeed0 = bt.feeds.PandasData(dataname=data0)
    cerebro.adddata(datafeed0, name=f'{index}')

    date_data = set_date()
    st_dt = date_data['st_dt']
    et_dt = date_data['et_dt']

    for i, code in enumerate(data_all):
        if code == index:
            df = data_all[code]
            df = df.loc[:, ['close']]
            df.rename(columns={'close': '399300'}, inplace=True)
            continue
        data = data_all[code]
        if len(data) == 0:
            continue
        print(f"i={i}->len={len(data)}->{code} Done !")
        data_bf = data.loc[data['datetime'] < st_dt]
        data = data.loc[data['datetime'] >= st_dt]

        if len(data_bf) > 0:
            info_h = data_bf.loc[data_bf['high'].idxmax(), :]
            info_l = data_bf.loc[data_bf['low'].idxmin(), :]
            stocks_info[code] = {'hdt': info_h['datetime'], 'high': info_h['high'], 'ldt': info_l['datetime'],
                                 'low': info_l['low']}

        datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_dt, todate=et_dt)
        cerebro.adddata(datafeed, name=code)
        print(f'{code} feeds ok')

    # 初始资金 1,000,000.0
    cerebro.broker.setcash(init_cash)
    # 佣金，双边各 0.0003
    cerebro.broker.setcommission(commission=0.0003)
    # 滑点：双边各 0.0001
    cerebro.broker.set_slippage_perc(perc=0.0001)
    # cerebro.addsizer(bt.sizers.FixedSize, stake=100)

    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='pnl')  # 返回收益率时序数据
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='_AnnualReturn')  # 年化收益率
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='_SharpeRatio')  # 夏普比率
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='_DrawDown')  # 回撤
    cerebro.addanalyzer(bt.analyzers.PyFolio, _name='pyfolio')

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    result = cerebro.run()
    strat = result[0]
    # st = strat.params.st
    pyfoliozer = strat.analyzers.getbyname('pyfolio')
    # 返回日度收益率序列
    daily_return = pd.Series(strat.analyzers.pnl.get_analysis())
    # 打印评价指标
    print("--------------- TimeReturn -----------------")
    # print(daily_return)
    print("--------------- AnnualReturn -----------------")
    print(strat.analyzers._AnnualReturn.get_analysis())
    print("--------------- SharpeRatio -----------------")
    print(strat.analyzers._SharpeRatio.get_analysis())
    print("--------------- DrawDown -----------------")
    print(strat.analyzers._DrawDown.get_analysis())
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())

    returns, positions, transactions, gross_lev = pyfoliozer.get_pf_items()
    # print(st_dt.date())
    returns = returns[returns.index.date >= st_dt.date()]
    returns.index = returns.index.date
    returns.index = pd.to_datetime(returns.index)

    # print(f'returns={returns}')
    # html_file = 'test.html'
    # qs.reports.html(returns, benchmark=df, title='grid', output=html_file, download_filename=html_file)
    # qs.reports.basic(returns)
    # f = webbrowser.open(html_file)
    # pf.create_full_tear_sheet(returns)
    # cerebro.plot()


class GridReversionStrategy(bt.Strategy):
    params = (
        ('period', 25),
        ('money', 3e4),
        ('ratio', 0.02),
        ('printlog', False),
        ('maperiod', 26),
        ('times', 1),
        ('span', 220),
        ('style', 3),
        ('st', datetime(2023, 1, 1)),
    )

    def __init__(self):
        self.order = None
        self.daily_pnl = dict()
        self.trade_info = list()
        self.pos_info = dict()
        self.trail_dict = dict()
        self.buy_first = dict()  # 每支ETF是否开始第一天就跌破之前最高价的buy_ratio
        self.cache_log = list()
        self.base_price = dict()  # 每个标的的基准价

        # print(f'stocks_info={stocks_info}')

        stocks = list()
        for data in self.datas:
            code = data._name
            if code != '399300':
                stocks.append(code)
                # print(f'stocks_info_in={stocks_info}')
                self.trail_dict[code] = {'hdt': stocks_info[code]['hdt'], 'high': stocks_info[code]['high'],
                                         'ldt': stocks_info[code]['ldt'], 'low': stocks_info[code]['low'],
                                         'costs': dict(), 'pos': dict(), 'count': 0, 'last_price_index': None,
                                         'amounts': dict(), 'profit': 0}

    def downcast(self, price):
        amount = self.p.money / price // 100 * 100
        return amount

    def set_grid(self, code, price):
        perc_level = [x for x in
                      np.arange(1 + self.p.ratio * 5, 1 - self.p.ratio * 5 - self.p.ratio / 2, -self.p.ratio)]
        price_levels = [price * x for x in perc_level]
        self.log(f'price={price}')
        self.log(f'perc_level={perc_level}')
        self.log(f'price_levels={price_levels}')
        self.trail_dict[code]['grids'] = price_levels
        self.trail_dict[code]['mid'] = len(price_levels) // 2
        self.log(50 * '-')

    def prenext(self):
        self.next()

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.datas[0].datetime.date(0)
            content = f'{dt.isoformat()}, {txt}'
            self.cache_log.append(content)
            print(content)

    @staticmethod
    def percent(today, yesterday):
        return float(today - yesterday) / today

    def next(self):
        if self.order:
            return

        cur_dt = self.datas[0].datetime.datetime(0)
        dt_int = int(cur_dt.strftime('%Y%m%d'))

        if dt_int not in self.daily_pnl:
            self.daily_pnl[dt_int] = round(self.broker.getvalue() - init_cash, 2)

        for data in self.datas:
            code = data._name
            try:
                dt = data.datetime.datetime(0)  # 获取当前的回测时间点
            except:
                print(f'{code}-{dt}')
            # dt != cur_dt 防止停牌日触发买卖
            if code == '399300' or dt != cur_dt:
                continue
            elif dt < self.p.st:
                continue

            # 开仓
            if self.trail_dict[code]['last_price_index'] is None:
                self.set_grid(code, data.close[-1])
                mid = self.trail_dict[code]['mid']
                grids = self.trail_dict[code]['grids']
                self.trail_dict[code]['last_price_index'] = mid
                total_amount = 0
                self.trail_dict[code]['costs'][mid] = grids[mid]
                for i in range(mid + 1, len(grids)):
                    if data.close[0] < grids[i]:
                        self.trail_dict[code]['last_price_index'] = i
                        self.trail_dict[code]['costs'][i] = grids[i]
                        amount = self.downcast(data.close[0])
                        self.trail_dict[code]['amounts'][i] = amount
                        total_amount += amount
                        self.log(f'开仓->amount={amount}->cost={grids[i]}->close={data.close[0]}')
                if total_amount > 0:
                    self.order = self.buy(data=data, size=total_amount, name=code)
            # 调仓
            else:
                buy_signal = False
                sell_signal = False
                total_amount = 0
                mid = self.trail_dict[code]['mid']
                grids = self.trail_dict[code]['grids']
                while True:
                    upper = None
                    lower = None
                    if self.trail_dict[code]['last_price_index'] > mid:
                        upper = grids[self.trail_dict[code]['last_price_index'] - 1]
                    if self.trail_dict[code]['last_price_index'] < len(grids) - 1:
                        lower = grids[self.trail_dict[code]['last_price_index'] + 1]
                    # 还不是最轻仓，继续涨，再卖一档
                    if upper is not None and data.close[0] > upper:
                        last_price_index = self.trail_dict[code]['last_price_index']
                        upper_index = last_price_index - 1
                        amount = self.trail_dict[code]['amounts'][last_price_index]
                        total_amount += amount
                        print(f'upper_index={upper_index}->last_price_index={last_price_index}')
                        if self.p.style == 1:
                            value = 0.022
                            r = random.uniform(-value, value)
                            profit = self.trail_dict[code]['costs'][last_price_index] * self.p.ratio * (1 + r) * amount
                        elif self.p.style == 2:
                            value = 0.018
                            r = random.uniform(-value, value)
                            profit = self.trail_dict[code]['costs'][last_price_index] * self.p.ratio * (1 + r) * amount
                        else:
                            value = 0.02
                            r = random.uniform(-value, value)

                            profit = self.trail_dict[code]['costs'][last_price_index] * self.p.ratio * (1 + r) * amount

                        self.trail_dict[code]['profit'] += profit
                        self.log(f'code={code}->sig_profit={profit}')
                        del self.trail_dict[code]['amounts'][last_price_index]
                        del self.trail_dict[code]['costs'][last_price_index]
                        self.trail_dict[code]['last_price_index'] = upper_index
                        sell_signal = True
                        self.set_grid(code, data.close[0])
                        self.log(f'卖出->amount={amount}->profit={profit}->close={data.close[0]}')
                        continue
                    # 还不是最重仓，继续跌，再买一档
                    if lower is not None and data.close[0] < lower:
                        last_price_index = self.trail_dict[code]['last_price_index']
                        self.trail_dict[code]['last_price_index'] = last_price_index + 1
                        self.trail_dict[code]['costs'][last_price_index + 1] = grids[last_price_index + 1]
                        amount = self.downcast(data.close[0])
                        self.log(
                            f"买入->amount={amount}->cost={self.trail_dict[code]['costs'][last_price_index + 1]}->close={data.close[0]}")
                        self.trail_dict[code]['amounts'][last_price_index + 1] = amount
                        total_amount += amount
                        buy_signal = True
                        self.set_grid(code, data.close[0])
                        continue
                    break
                if buy_signal:
                    self.order = self.buy(data=data, size=total_amount, name=code)
                elif sell_signal:
                    self.order = self.sell(data=data, size=total_amount, name=code)

    def notify_order(self, order):
        code = order.data._name

        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:

            if order.isbuy():
                self.log(
                    f'code={code} BUY CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={code} BUY EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f'Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')
            else:
                self.log(
                    f'code={order.data._name} SELL CREATE TIME: {bt.num2date(order.created.dt)}, EXECUTED TIME: {bt.num2date(order.executed.dt)}')
                self.log(
                    f'code={order.data._name} SELL EXECUTED, ref:{order.ref}，Price: {order.executed.price:.2f}, Cost: {order.executed.value: .2f},'
                    f' Comm {order.executed.comm:.2f}, Size: {order.executed.size:.2f}')

            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log(f'Order Canceled/Margin/Rejected->code={code}')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        self.log(f'OPERATION PROFIT, GROSS {trade.pnl:.2f}, NET {trade.pnlcomm:.2f}')

    def stop(self):
        print('Strategy Finish!!!')
        res = list()
        for code in self.trail_dict:
            profit = self.trail_dict[code]['profit']
            # if profit < 1 or profit > 20000:
            #     continue
            res.append([code, profit])

        data = pd.DataFrame(res, columns=['code', 'profit'])
        print(data)
        data.to_excel(r'C:\Users\AndrewX\Desktop\下降_3.xlsx', sheet_name='模式C')

        # with open(r'C:\Users\AndrewX\Desktop\cache_log.txt', 'w') as f:
        #     res = '\n'.join(self.cache_log)
        #     f.write(res)
        doprint = (False, True)[0]

        if doprint:
            daily_pnl_df = pd.DataFrame(self.daily_pnl.items(), columns=['DATE_T', 'PNL'])
            append_df = pd.DataFrame([[101, round((self.broker.getvalue() / init_cash - 1) * 100, 2)]],
                                     columns=['DATE_T', 'PNL'])
            daily_pnl_df = daily_pnl_df.append(append_df, ignore_index=True)
            trade_df = pd.DataFrame(self.trade_info,
                                    columns=['DATE_T', 'STK_CODE', 'ACTION', 'PRICE', 'AMOUNT', 'COMM'])
            pnl_table = 'TEST_DAILY_PNL'
            trade_table = 'TEST_TRADE_INFO'
            # daily_pnl_df = daily_pnl_df.loc[(daily_pnl_df['DATE_T'] >= 20230428) | (daily_pnl_df['DATE_T'] == 101)]
            # insert_db(daily_pnl_df, pnl_table)
            # insert_db(trade_df, trade_table)
            print(self.trail_dict)
            print('insert table ok')


def set_date():
    bt = "20220101"
    st = "20230101"
    et = "20230816"
    bt_dt = datetime.strptime(st, '%Y%m%d')
    st_dt = datetime.strptime(st, '%Y%m%d')
    et_dt = datetime.strptime(et, '%Y%m%d')
    bt_d = bt_dt.date()
    st_d = st_dt.date()
    et_d = et_dt.date()
    res = {'bt': bt, 'st': st, 'et': et, 'bt_d': bt_d, 'st_d': st_d, 'et_d': et_d, 'bt_dt': bt_dt, 'st_dt': st_dt,
           'et_dt': et_dt}

    return res


def get_excel_stocks():
    data = pd.read_excel(r'C:\Users\AndrewX\Desktop\标的 20231012.xlsx', header=None, sheet_name='下降')
    stocks = data.iloc[:, 0].to_list()
    stocks = [tok[-6:] for tok in stocks]
    return stocks


def main():
    index = '399300'
    # stocks = get_stocks()
    # stocks = ['510880', '159915', '513100', '518880']
    # stocks = ['159901', '510880']
    # stocks = ['159901', '510880', '512510', '512220', '159938', '159909', '513500', '513880', '515770', '515350',
    #           '159810', '517350']
    # stocks = stocks[:200]
    # stocks = ['159909', '510170', '159915', '512220', '510880', '510300']
    #
    # stocks = [('159918', '159909', '512220', '159915', '510500', '510510', '159907', '159922', '510290', '159936',
    #            '159930', '159935', '159912', '159903', '159822', '513050', '159902', '510900', '512070', '510410')]
    # stocks = ['513180']
    # print(stocks)
    stocks = get_excel_stocks()
    # stocks = stocks[:500]
    # stocks = [str(tok) for tok in stocks]
    # stocks = ['159721', '159733', '159735', '159745', '159766', '159773', '159783',
    #           '159790', '159792', '159808']
    stocks.insert(0, index)
    run_backtrade(stocks)
    # s = get_data(stocks, freq='day', source='not db')
    # print(s)


if __name__ == '__main__':
    main()
